#!/usr/bin/python

import os, re, subprocess
from common import *
import globals

timeCodeFiles = []

def cleanup(streams, logname):
    for s in streams:
        streamdata = streams[s]
        if streamdata.has_key('fname'):
            removeFile(streamdata['fname'])
        if streamdata.has_key('IFO'):
            removeFile(streamdata['IFO'])
        if streamdata.has_key('idx'):
            removeFile(streamdata['idx'])
        if streamdata.has_key('sub'):
            removeFile(streamdata['sub'])
    removeFile('chapters.xml')
    while timeCodeFiles:
        removeFile(timeCodeFiles.pop())

def getLogName(mkvfile):
    return os.path.splitext(os.path.split(mkvfile)[1])[0] + '_log.txt'

def getstreams(mkvfile):
    output = subprocess.Popen(['mkvinfo', mkvfile], stdout = subprocess.PIPE).communicate()[0]
    print output

    tracks = {}
    trackNumber = None
    for line in output.splitlines():
        r = re.search('[+] Track number: (\d+)', line)
        if r:
            id = re.search('track ID for mkvmerge & mkvextract: (\d+)', line)
            if id:
                trackNumber = int(id.group(1))
            else:
                trackNumber = int(r.group(1))
            continue

        r = re.search('[+] Track type: (.+)', line)
        if r:
            trackType = r.group(1)
            if tracks.has_key(trackNumber):
                tracks[trackNumber]['type'] = trackType
            else:
                tracks[trackNumber] = {'type': trackType}
            continue

        r = re.search('[+] Language: (.+)', line)
        if r:
            language = r.group(1)
            if tracks.has_key(trackNumber):
                tracks[trackNumber]['language'] = language
            else:
                tracks[trackNumber] = {'language': language}
            continue
    return tracks

def resolvefilenames(tracks, logfile):
    #cont = open(logfile).read()
    for trackNumber in tracks:
        fname = 'track_%d_%s' % (trackNumber, tracks[trackNumber]['type'])
        tracks[trackNumber]['fname'] = fname
        if tracks[trackNumber]['type'] in ['subtitles'] :
            tracks[trackNumber]['sub'] = fname + '.sub'
            tracks[trackNumber]['idx'] = fname + '.idx'
            #tracks[trackNumber]['IFO'] = fname + '.IFO'
    return tracks

def demux(mkvfile):
    fname = mkvfile

    baseName = os.path.splitext(os.path.split(fname)[1])[0]
    logName = getLogName(mkvfile)
    tracks = resolvefilenames(getstreams(mkvfile), logName)

    print "fname =", fname
    #print "logName =", logName
    print "demuxing..."
    execute(r'mkvextract chapters "%s" >chapters.xml' % (mkvfile))

    cmd = 'mkvextract timecodes_v2 "%s"' % (mkvfile)
    for trackNumber in tracks:
        timeCodeFile = 'tc-track%d.txt' % (trackNumber)
        timeCodeFiles.append(timeCodeFile)
        cmd += ' %d:%s' % (trackNumber, timeCodeFile)
    execute(cmd)

    cmd = 'mkvextract tracks "%s" ' % (mkvfile)
    for trackNumber in tracks:
        cmd += '%d:%s ' % (trackNumber, tracks[trackNumber]['fname'])
    execute(cmd)
    print "done"

if __name__ == '__main__':
    import sys

    mkvfile = sys.argv[1]
    logName = getLogName(mkvfile)

    demux(mkvfile)

    print 'SUMMARY OF STREAMS:'
    streams = resolvefilenames(getstreams(mkvfile), logName)
    for s in streams:
        print 'Track:', s, streams[s]
